package com.yhkj.data.datasource;

import com.yhkj.data.mq.message.DataSourceConfig;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Component;

import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.function.Supplier;

/**
 * @author <a href="mailto:idler41@163.con">linfuxin</a> created on 2023-03-08 10:19:37
 */
@Component
@Slf4j
public class JdbcManager {

    @Autowired
    private RoutingDataSource routingDataSource;
    private static final int EXECUTE_BATCH_SIZE = 200;

    public List<Object[]> executeQueryPkList(DataSourceConfig dataSourceConfig, String sql, int keysSize, List<Object> params) {
        return switchAndExecute(dataSourceConfig, () -> {
            try (
                    Connection connection = routingDataSource.getConnection();
                    PreparedStatement ps = connection.prepareStatement(sql);
            ) {
                for (int i = 0, size = params.size(); i < size; i++) {
                    ps.setObject(i + 1, params.get(i));
                }
                List<Object[]> result = new ArrayList<>(params.size());
                try (ResultSet rs = ps.executeQuery()) {
                    while (rs.next()) {
                        Object[] keyDataRow = new Object[keysSize];
                        for (int r = 0; r < keysSize; r++) {
                            keyDataRow[r] = rs.getObject(r + 1);
                        }
                        result.add(keyDataRow);
                    }
                    return result;
                }
            } catch (SQLException e) {
                log.error("执行sql异常!sql={},params={}", sql, params);
                throw new RuntimeException(e);
            }
        });
    }

    public Integer executeBatch(DataSourceConfig dataSourceConfig, String sql, List<List<Object>> params) {
        return switchAndExecute(dataSourceConfig, () -> {
            try (
                    Connection connection = routingDataSource.getConnection();
                    PreparedStatement ps = connection.prepareStatement(sql);
            ) {
                int rowAffected = 0;
                for (int i = 0, size = params.size(); i < size; i++) {
                    List<Object> rowParam = params.get(i);
                    for (int j = 0, cSize = rowParam.size(); j < cSize; j++) {
                        // ps下标从1开始
                        ps.setObject(j + 1, rowParam.get(j));
                    }
                    ps.addBatch();
                    if (i % EXECUTE_BATCH_SIZE == 0) {
                        int[] batchResults = ps.executeBatch();
                        rowAffected += Arrays.stream(batchResults).sum();
                        ps.clearBatch();
                    }
                }
                int[] batchResults = ps.executeBatch();
                rowAffected += Arrays.stream(batchResults).sum();
                ps.clearBatch();
                return rowAffected;
            } catch (SQLException e) {
                log.error("执行sql异常!sql={},params={}", sql, params);
                throw new RuntimeException(e);
            }
        });
    }


    public <V> V switchAndExecute(DataSourceConfig dataSourceConfig, Supplier<V> supplier) {
        try {
            LocalRoutingKeyHolder.setContextKey(dataSourceConfig);
            if (log.isDebugEnabled()) {
                log.debug("switch dataSource!key={}", dataSourceConfig.getJdbcUrl());
            }
            return supplier.get();
        } finally {
            LocalRoutingKeyHolder.removeContextKey();
        }
    }
}
